import numpy as np
import cv2
import onnxruntime as ort
from onnx_utils_special import nms
import os
import matplotlib.pyplot as plt
from yolo_decode import *


def sort_keypoints(keypoints):
    """
    Sort the keypoints by their position in the image.

    @param keypoints: 4行2列的关键点坐标，numpy数组，每行代表一个关键点的坐标(x,y)
    @return: 按左上、右上、右下、左下顺序排列的关键点坐标
    """
    index_left_up = np.argmin(np.sum(keypoints * np.array([1, 1]), axis=1))
    index_left_down = np.argmin(np.sum(keypoints * np.array([1, -1]), axis=1))
    index_right_down = np.argmax(np.sum(keypoints * np.array([1, 1]), axis=1))
    index_right_up = np.argmax(np.sum(keypoints * np.array([1, -1]), axis=1))
    keypoints_sorted = np.array(
        [
            keypoints[index_left_up],
            keypoints[index_right_up],
            keypoints[index_right_down],
            keypoints[index_left_down],
        ]
    )
    return keypoints_sorted


if __name__ == "__main__":
    model = ort.InferenceSession(r"C:\Users\Administrator\Desktop\ultralytics-action\onnx_predict\last.onnx")
    # img = cv2.imread(r"D:\desktopD\aidc_local\models\keypoint_detection\gitignored_dir\x1test\bb1.png")
    save_dir = r"C:\Users\Administrator\Desktop\ultralytics-action\预测结果\special"
    os.makedirs(save_dir, exist_ok=True)

    img_directory_path = (
        r"D:\desktopD\aidc_local\models\keypoint_detection\app_runs_result\test_barcode_img_for_pose\2D"
    )
    images_name_list = os.listdir(img_directory_path)
    for img_name in images_name_list:
        if img_name.split(".")[-1] not in ["jpg", "png"]:
            continue
        img = cv2.imread(os.path.join(img_directory_path, img_name))
        img_rgb = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
        height, width, _ = img_rgb.shape
        length = max(height, width)
        imageB = np.zeros((256, 256, 3), np.float32)
        s = 256 / length

        resized_img = cv2.resize(img_rgb, (0, 0), fx=s, fy=s)
        h, w, c = resized_img.shape
        imageB[0:h, 0:w] = resized_img

        imageB = np.array([imageB.transpose(2, 0, 1)])
        imageB = np.mean(imageB, axis=1, keepdims=True)
        imageB /= 255

        output = model.run(None, {"input": imageB})
        out = output[0]

        detect_pose_data = out.transpose(0, 1, 2, 3)
        reg_max = 16
        nc = detect_pose_data.shape[1] - reg_max * 8
        detect_data = detect_pose_data[:, 0 : reg_max * 8 + nc, :].reshape(
            detect_pose_data.shape[0], reg_max * 8 + nc, -1
        )
        anchors, strides = (x for x in make_anchors([detect_pose_data], [16], 0.5))
        anchors = anchors.transpose(1, 0)
        strides = strides.transpose(1, 0)
        bbox = detect_data[:, 0 : reg_max * 8, :]
        cls = detect_data[:, reg_max * 8 :, :]
        dfled_bbox = dfl(bbox, point_num=4)
        # dbox = dist2bbox(dfled_bbox, anchors, xywh=True, dim=1) * strides / 256

        lt, rt, rb, lb = np.split(dfled_bbox, 4, axis=1)

        x1y1 = anchors + lt * np.array([-1, -1]).reshape(1, 2, 1)
        x1y1 = x1y1 * 16
        x2y2 = anchors + rt * np.array([1, -1]).reshape(1, 2, 1)
        x2y2 = x2y2 * 16
        x3y3 = anchors + rb * np.array([1, 1]).reshape(1, 2, 1)
        x3y3 = x3y3 * 16
        x4y4 = anchors + lb * np.array([-1, 1]).reshape(1, 2, 1)
        x4y4 = x4y4 * 16

        x_min = np.min(
            np.concatenate((x1y1[:, 0:1, :], x2y2[:, 0:1, :], x3y3[:, 0:1, :], x4y4[:, 0:1, :]), axis=1),
            axis=1,
            keepdims=True,
        )
        y_min = np.min(
            np.concatenate((x1y1[:, 1:2, :], x2y2[:, 1:2, :], x3y3[:, 1:2, :], x4y4[:, 1:2, :]), axis=1),
            axis=1,
            keepdims=True,
        )
        x_max = np.max(
            np.concatenate((x1y1[:, 0:1, :], x2y2[:, 0:1, :], x3y3[:, 0:1, :], x4y4[:, 0:1, :]), axis=1),
            axis=1,
            keepdims=True,
        )
        y_max = np.max(
            np.concatenate((x1y1[:, 1:2, :], x2y2[:, 1:2, :], x3y3[:, 1:2, :], x4y4[:, 1:2, :]), axis=1),
            axis=1,
            keepdims=True,
        )
        bbox = np.concatenate(((x_min + x_max) / 2, (y_min + y_max) / 2, (x_max - x_min), (y_max - y_min)), axis=1)
        cls = sigmoid(cls)
        detect_data = np.concatenate((bbox, cls, x1y1, x2y2, x3y3, x4y4), axis=1)
        res = nms(prediction=detect_data)[0]
        colors1 = [255, 0, 0]
        # 矩形框的颜色
        colors2 = [0, 255, 255]
        # 关键点的颜色
        colors3 = [255, 255, 0]
        if res is not None:
            for rect in res:
                rect = list(map(int, rect / 256 * length))
                cv2.rectangle(
                    img_rgb,
                    (rect[0] - rect[2] // 2, rect[1] - rect[3] // 2),
                    (rect[0] + rect[2] // 2, rect[1] + rect[3] // 2),
                    colors1,
                    3,
                )
                polylines_array = np.array(
                    [
                        [rect[6], rect[7]],
                        [rect[8], rect[9]],
                        [rect[10], rect[11]],
                        [rect[12], rect[13]],
                    ]
                ).astype(np.int32)
                polylines_array = sort_keypoints(polylines_array)
                cv2.polylines(img_rgb, [polylines_array], True, colors2, 3)

                r = img.shape[0] // 100
                cv2.circle(img_rgb, (int(rect[6]), int(rect[7])), r, colors3, -1)
                cv2.circle(img_rgb, (int(rect[8]), int(rect[9])), r, colors3, -1)
                cv2.circle(img_rgb, (int(rect[10]), int(rect[11])), r, colors3, -1)
                cv2.circle(img_rgb, (int(rect[12]), int(rect[13])), r, colors3, -1)

        plt.imsave(os.path.join(save_dir, img_name), img_rgb)
